#include <NvInfer.h>
#include <cuda_runtime_api.h>

#include <opencv2/opencv.hpp>

#include <algorithm>
#include <chrono>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <memory>
#include <numeric>
#include <string>
#include <vector>

namespace fs = std::filesystem;

class TrtLogger : public nvinfer1::ILogger {
public:
    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override {
        if (severity <= nvinfer1::ILogger::Severity::kWARNING) {
            std::cout << "[TRT] " << msg << std::endl;
        }
    }
};

static std::vector<unsigned char> readFile(const std::string& path) {
    std::ifstream ifs(path, std::ios::binary | std::ios::ate);
    if (!ifs) {
        throw std::runtime_error("Failed to open file: " + path);
    }
    std::streamsize size = ifs.tellg();
    ifs.seekg(0, std::ios::beg);
    std::vector<unsigned char> buffer(size);
    if (!ifs.read(reinterpret_cast<char*>(buffer.data()), size)) {
        throw std::runtime_error("Failed to read file: " + path);
    }
    return buffer;
}

static void checkCuda(cudaError_t e, const char* msg) {
    if (e != cudaSuccess) {
        throw std::runtime_error(std::string("CUDA Error: ") + msg + ": " + cudaGetErrorString(e));
    }
}

static cv::Mat colorizeMask(const cv::Mat& mask) {
    // Cityscapes colormap (19 classes)
    static const cv::Vec3b colors[19] = {
        {128, 64, 128}, {244, 35, 232}, {70, 70, 70}, {102, 102, 156}, {190, 153, 153},
        {153, 153, 153}, {250, 170, 30}, {220, 220, 0}, {107, 142, 35}, {152, 251, 152},
        {70, 130, 180}, {220, 20, 60}, {255, 0, 0}, {0, 0, 142}, {0, 0, 70},
        {0, 60, 100}, {0, 80, 100}, {0, 0, 230}, {119, 11, 32}
    };

    cv::Mat color(mask.size(), CV_8UC3);
    for (int y = 0; y < mask.rows; ++y) {
        const uint8_t* pm = mask.ptr<uint8_t>(y);
        cv::Vec3b* pc = color.ptr<cv::Vec3b>(y);
        for (int x = 0; x < mask.cols; ++x) {
            uint8_t id = pm[x];
            if (id < 19) pc[x] = colors[id];
            else pc[x] = {0, 0, 0};
        }
    }
    return color;
}

int main(int argc, char** argv) {
    std::string enginePath = "checkpoints/seg.trt";
    std::string imagePath = "samples/CAM_FRONT.jpg";
    std::string outputDir = "results";

    // Parse simple CLI
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        auto next = [&](int& i) -> std::string {
            if (i + 1 >= argc) throw std::runtime_error("Missing value for argument: " + a);
            return std::string(argv[++i]);
        };
        if (a == "--engine") enginePath = next(i);
        else if (a == "--input") imagePath = next(i);
        else if (a == "--outdir") outputDir = next(i);
        else if (a == "-h" || a == "--help") {
            std::cout << "Usage: " << argv[0] << " [--engine checkpoints/seg.trt] [--input samples/CAM_FRONT.jpg] [--outdir results]\n";
            return 0;
        }
    }

    const int targetH = 1024;
    const int targetW = 2048;
    const int inputC = 3;
    const int numClasses = 19;

    try {
        fs::create_directories(outputDir);

        // Load image
        cv::Mat originalBGR = cv::imread(imagePath, cv::IMREAD_COLOR);
        if (originalBGR.empty()) {
            std::cerr << "Failed to read image: " << imagePath << std::endl;
            return 1;
        }
        const int origH = originalBGR.rows;
        const int origW = originalBGR.cols;

        // Preprocess: resize to (W,H), BGR->RGB, [0,1], normalize by ImageNet mean/std, CHW
        auto t0 = std::chrono::steady_clock::now();
        cv::Mat resizedBGR;
        cv::resize(originalBGR, resizedBGR, cv::Size(targetW, targetH), 0, 0, cv::INTER_LINEAR);
        cv::Mat rgb;
        cv::cvtColor(resizedBGR, rgb, cv::COLOR_BGR2RGB);
        cv::Mat rgbF;
        rgb.convertTo(rgbF, CV_32FC3, 1.0 / 255.0);
        const float mean[3] = {0.485f, 0.456f, 0.406f};
        const float stdv[3] = {0.229f, 0.224f, 0.225f};

        std::vector<float> inputHost(inputC * targetH * targetW);
        for (int y = 0; y < targetH; ++y) {
            const cv::Vec3f* pr = rgbF.ptr<cv::Vec3f>(y);
            for (int x = 0; x < targetW; ++x) {
                for (int c = 0; c < inputC; ++c) {
                    float v = (pr[x][c] - mean[c]) / stdv[c];
                    inputHost[c * targetH * targetW + y * targetW + x] = v;
                }
            }
        }
        auto t1 = std::chrono::steady_clock::now();

        // TensorRT runtime and engine
        TrtLogger logger;
        auto engineData = readFile(enginePath);
        std::unique_ptr<nvinfer1::IRuntime, void(*)(nvinfer1::IRuntime*)> runtime(
            nvinfer1::createInferRuntime(logger), [](nvinfer1::IRuntime* p){ if(p) p->destroy(); }
        );
        if (!runtime) throw std::runtime_error("Failed to create TensorRT runtime");

        std::unique_ptr<nvinfer1::ICudaEngine, void(*)(nvinfer1::ICudaEngine*)> engine(
            runtime->deserializeCudaEngine(engineData.data(), engineData.size()),
            [](nvinfer1::ICudaEngine* p){ if(p) p->destroy(); }
        );
        if (!engine) throw std::runtime_error("Failed to deserialize engine: " + enginePath);

        std::unique_ptr<nvinfer1::IExecutionContext, void(*)(nvinfer1::IExecutionContext*)> context(
            engine->createExecutionContext(), [](nvinfer1::IExecutionContext* p){ if(p) p->destroy(); }
        );
        if (!context) throw std::runtime_error("Failed to create execution context");

        // Resolve bindings (assume single input and single output)
        int nbBindings = engine->getNbBindings();
        int inputIndex = -1, outputIndex = -1;
        for (int i = 0; i < nbBindings; ++i) {
            if (engine->bindingIsInput(i)) inputIndex = i; else outputIndex = i;
        }
        if (inputIndex < 0 || outputIndex < 0) throw std::runtime_error("Invalid engine bindings");

        // Set binding dims for dynamic shapes if needed
        nvinfer1::Dims inputDims = engine->getBindingDimensions(inputIndex);
        if (inputDims.nbDims == 3) {
            // Expect CHW with explicit batch; set N via optimization profile (implicit handled by TRT)
            nvinfer1::Dims4 d(1, inputC, targetH, targetW);
            context->setBindingDimensions(inputIndex, d);
        } else if (inputDims.nbDims == 4) {
            // Potentially NCHW; ensure dimensions are set
            nvinfer1::Dims4 d(1, inputC, targetH, targetW);
            context->setBindingDimensions(inputIndex, d);
        }

        // Allocate device buffers
        size_t inputCount = static_cast<size_t>(1) * inputC * targetH * targetW;
        size_t outputCount = static_cast<size_t>(1) * numClasses * targetH * targetW;
        void* dInput = nullptr;
        void* dOutput = nullptr;
        checkCuda(cudaMalloc(&dInput, inputCount * sizeof(float)), "cudaMalloc dInput");
        checkCuda(cudaMalloc(&dOutput, outputCount * sizeof(float)), "cudaMalloc dOutput");

        void* bindings[2];
        bindings[inputIndex] = dInput;
        bindings[outputIndex] = dOutput;

        // Copy input to device
        checkCuda(cudaMemcpy(dInput, inputHost.data(), inputCount * sizeof(float), cudaMemcpyHostToDevice), "cudaMemcpy H2D");

        // Inference
        auto t2 = std::chrono::steady_clock::now();
        if (!context->enqueueV2(bindings, nullptr, nullptr)) {
            throw std::runtime_error("Inference enqueueV2 failed");
        }
        checkCuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize");
        auto t3 = std::chrono::steady_clock::now();

        // Copy output back
        std::vector<float> outputHost(outputCount);
        checkCuda(cudaMemcpy(outputHost.data(), dOutput, outputCount * sizeof(float), cudaMemcpyDeviceToHost), "cudaMemcpy D2H");

        // Postprocess: argmax over classes -> mask, colorize, resize back, overlay
        auto t4 = std::chrono::steady_clock::now();
        cv::Mat mask(targetH, targetW, CV_8UC1);
        for (int y = 0; y < targetH; ++y) {
            uint8_t* pm = mask.ptr<uint8_t>(y);
            for (int x = 0; x < targetW; ++x) {
                int idxBase = y * targetW + x;
                int bestId = 0;
                float bestVal = outputHost[idxBase]; // class 0
                for (int c = 1; c < numClasses; ++c) {
                    float v = outputHost[c * targetH * targetW + idxBase];
                    if (v > bestVal) { bestVal = v; bestId = c; }
                }
                pm[x] = static_cast<uint8_t>(bestId);
            }
        }
        cv::Mat colorRGB = colorizeMask(mask); // currently in RGB as per table
        // Resize to original size
        cv::Mat colorRGBOrig;
        cv::resize(colorRGB, colorRGBOrig, cv::Size(origW, origH), 0, 0, cv::INTER_NEAREST);
        // Overlay
        cv::Mat colorBGROrig;
        cv::cvtColor(colorRGBOrig, colorBGROrig, cv::COLOR_RGB2BGR);
        cv::Mat overlay;
        const double alpha = 0.7; // same as Python
        cv::addWeighted(originalBGR, 1.0 - alpha, colorBGROrig, alpha, 0.0, overlay);
        auto t5 = std::chrono::steady_clock::now();

        // Save results
        auto base = fs::path(imagePath).stem().string();
        std::string segPath = (fs::path(outputDir) / (base + std::string("_segmentation.png"))).string();
        std::string overlayPath = (fs::path(outputDir) / (base + std::string("_overlay.png"))).string();
        std::string maskPath = (fs::path(outputDir) / (base + std::string("_mask.png"))).string();

        cv::imwrite(segPath, colorBGROrig);
        cv::imwrite(overlayPath, overlay);
        cv::Mat maskOrig;
        cv::resize(mask, maskOrig, cv::Size(origW, origH), 0, 0, cv::INTER_NEAREST);
        cv::imwrite(maskPath, maskOrig);

        // Timing
        auto durPrep = std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
        auto durInf = std::chrono::duration_cast<std::chrono::milliseconds>(t3 - t2).count();
        auto durPost = std::chrono::duration_cast<std::chrono::milliseconds>(t5 - t4).count();

        std::cout << "Preprocessing time: " << durPrep / 1000.0 << "s\n";
        std::cout << "Inference time: " << durInf / 1000.0 << "s\n";
        if (durInf > 0) {
            std::cout << "FPS: " << 1000.0 / durInf << "\n";
        }
        std::cout << "Postprocessing time: " << durPost / 1000.0 << "s\n";
        std::cout << "Segmentation result saved to: " << segPath << "\n";
        std::cout << "Overlay result saved to: " << overlayPath << "\n";
        std::cout << "Prediction mask saved to: " << maskPath << "\n";

        // Cleanup
        cudaFree(dInput);
        cudaFree(dOutput);

        return 0;
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
} 